import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

# -------------------------- 1. 3D体素数据表示 --------------------------
# 体素（Voxel）：3D网格，每个格子表示“是否存在物体”（1=存在，0=不存在）
def create_cube_voxel(size=32):
    """生成立方体的3D体素数据"""
    voxel = torch.zeros(size, size, size)
    # 立方体范围：中间区域
    voxel[10:22, 10:22, 10:22] = 1.0
    return voxel

def create_sphere_voxel(size=32):
    """生成球体的3D体素数据"""
    voxel = torch.zeros(size, size, size)
    center = size // 2
    radius = 10
    # 遍历3D网格，判断是否在球内
    for x in range(size):
        for y in range(size):
            for z in range(size):
                if (x-center)**2 + (y-center)**2 + (z-center)**2 <= radius**2:
                    voxel[x, y, z] = 1.0
    return voxel


# -------------------------- 2. 3D模型可视化 --------------------------
def plot_voxel(voxel, title="3D Voxel Model"):
    """可视化3D体素模型"""
    fig = plt.figure(figsize=(8, 8))
    ax = fig.add_subplot(111, projection='3d')
    # 提取体素中“存在物体”的坐标
    x, y, z = np.where(voxel.numpy() > 0.5)
    ax.scatter(x, y, z, c='blue', marker='s', s=10)
    ax.set_title(title)
    plt.show()


# -------------------------- 3. 用PyTorch 3D网络生成3D模型 --------------------------
class Simple3DGenerator(nn.Module):
    """简单的3D生成网络：输入随机噪声，输出32×32×32体素"""
    def __init__(self, latent_dim=100):
        super(Simple3DGenerator, self).__init__()
        self.latent_dim = latent_dim
        # 3D转置卷积（上采样）：从噪声生成3D体素
        self.model = nn.Sequential(
            # 输入：latent_dim → 转为4×4×4的特征图
            nn.ConvTranspose3d(latent_dim, 512, kernel_size=4, stride=1, padding=0),
            nn.ReLU(),
            # 上采样到8×8×8
            nn.ConvTranspose3d(512, 256, kernel_size=4, stride=2, padding=1),
            nn.ReLU(),
            # 上采样到16×16×16
            nn.ConvTranspose3d(256, 128, kernel_size=4, stride=2, padding=1),
            nn.ReLU(),
            # 上采样到32×32×32，输出单通道体素（1=存在，0=不存在）
            nn.ConvTranspose3d(128, 1, kernel_size=4, stride=2, padding=1),
            nn.Sigmoid()  # 归一化到[0,1]
        )

    def forward(self, z):
        # 输入噪声：(batch_size, latent_dim) → 转为(batch_size, latent_dim, 1, 1, 1)
        z = z.view(z.size(0), self.latent_dim, 1, 1, 1)
        voxel = self.model(z)
        return voxel.squeeze(1)  # 去掉通道维度，输出(batch_size, 32, 32, 32)


# -------------------------- 运行代码 --------------------------
if __name__ == "__main__":
    # 1. 生成基础3D体素并可视化
    cube_voxel = create_cube_voxel()
    plot_voxel(cube_voxel, title="Cube Voxel Model")

    sphere_voxel = create_sphere_voxel()
    plot_voxel(sphere_voxel, title="Sphere Voxel Model")

    # 2. 用生成网络生成3D模型
    generator = Simple3DGenerator(latent_dim=100)
    # 输入随机噪声
    noise = torch.randn(1, 100)  # batch_size=1
    generated_voxel = generator(noise)
    # 二值化（0.5为阈值）
    generated_voxel = (generated_voxel > 0.5).float()
    plot_voxel(generated_voxel, title="Generated 3D Voxel Model")